﻿using System;
using System.Linq;
using System.Runtime.InteropServices;
using System.Threading;
using Ast.Core.Events;
using Ast.Core.Services.Extensions;
using Ast.Framework;
using Castle.Core.Logging;

namespace Ast.Core.Services
{
    public interface IHookSystemService : ILockable, IDependency
    {
        void Hook();
    }

    public class HookSystemService : IHookSystemService, IDisposable
    {
        private readonly IKernelAddressFinder _kernelAddressFinder;
        private readonly IDeviceController _deviceController;
        private readonly IHookSystemEvents _hookSystemEvents;

        /*
            8192f6ab 2be1            sub     esp,ecx
            8192f6ad c1e902          shr     ecx,2
            8192f6ab  2b e1 c1 e9 02
        */
        private static readonly byte[] HookSignature = {0x2b, 0xe1, 0xc1, 0xe9, 0x02};

        private readonly object _syncRoot = new object();
        private IntPtr _addressOfHook = IntPtr.Zero;

        public HookSystemService(IKernelAddressFinder kernelAddressFinder, IDeviceController deviceController, IHookSystemEvents hookSystemEvents)
        {
            _kernelAddressFinder = kernelAddressFinder;
            _deviceController = deviceController;
            _hookSystemEvents = hookSystemEvents;

            Logger = NullLogger.Instance;
        }

        public ILogger Logger { get; set; }

        public void Hook()
        {
            Unhook();

            // get address of KiSystemService
            var idt = _kernelAddressFinder.FindIdt();
            var gdt = _kernelAddressFinder.FindGdt();

            var idt2E = _deviceController.ReadAs<Idt>(idt + 0x2e*8);
            Logger.InfoFormat("2e idt 0x{0:x}, selector 0x{1:x}, offset 0x{2:x}, type 0x{3:x}", idt2E.Value, idt2E.Selector, idt2E.Offset, idt2E.Type);
            var gdt2E = _deviceController.ReadAs<Gdt>(gdt + (int) (8*idt2E.SelectIndex));
            var adrOfKiSystemService = gdt2E.Base + idt2E.Offset;
            Logger.InfoFormat("2e gdt 0x{0:x}, base 0x{1:x}, KiSystemService address 0x{2:x}", gdt2E.Value, gdt2E.Base, adrOfKiSystemService);

            // search position to hook
            Logger.Info("searching hook signature");
            var contents = _deviceController.ReadBytes(new IntPtr((int) adrOfKiSystemService), 1024);
            var offset = contents.IndexOfSequence(HookSignature).FirstOrDefault();
            if (offset == 0)
            {
                Logger.ErrorFormat("hook address not found");
                return;
            }
            _addressOfHook = new IntPtr((int) (adrOfKiSystemService + offset));
            Logger.InfoFormat("found hook address 0x{0:x}, offset 0x{1:x}", _addressOfHook.ToInt32(), offset);

            // hook
            if (_deviceController.Send(IoCtl.IoCtlHooKInstall, _addressOfHook))
            {
                Logger.Info("hooked system");

                // events
                _hookSystemEvents.OnHooked();
            }
            else
            {
                Logger.Error("failed to hook system");
            }
        }

        private void Unhook(bool isDisposing = false)
        {
            AcquireLock();
            if (_addressOfHook != IntPtr.Zero)
            {
                _deviceController.Send(IoCtl.IoCtlHookUninstall, _addressOfHook);
                _addressOfHook = IntPtr.Zero;
                Logger.Info("unhooked system");

                // events
                if (!isDisposing) _hookSystemEvents.OnUnhooked();
            }
            ReleaseLock();
        }

        public void AcquireLock()
        {
            Monitor.Enter(_syncRoot);
        }

        public void ReleaseLock()
        {
            Monitor.Exit(_syncRoot);
        }

        public void Dispose()
        {
            Unhook(true);
        }
    }

    [StructLayout(LayoutKind.Sequential, CharSet = CharSet.Auto, Pack = 1)]
    public struct Idt
    {
        public ulong Value;

        public uint Selector
        {
            get { return (uint) ((Value >> 16) & 0xFFFF); }
        }

        public uint SelectIndex
        {
            get { return Selector >> 3; }
        }

        public uint Offset
        {
            get { return (uint) (((Value >> 48) << 16) + (Value & 0xFFFF)); }
        }

        public byte Type
        {
            get { return (byte) ((Value >> 40) & 0xF); }
        }
    }

    [StructLayout(LayoutKind.Sequential, CharSet = CharSet.Auto, Pack = 1)]
    public struct Gdt
    {
        public ulong Value;

        public uint Base
        {
            get { return (uint) (((Value >> 56) << 8) + ((Value >> 16) & 0xFFFFFF)); }
        }
    }
}